import os

from pysqlcipher3 import dbapi2 as sqlite


class EncryptDb():
    ''' 数据库操作 '''

    def __init__(self, source_dbs_path, db_file_name, decrypt_db_path, key, cipher_page_size=1024):
        self.key = key
        self.cipher_page_size = cipher_page_size
        self.db_file_name = db_file_name
        self.decrypt_db_path = decrypt_db_path
        self.conn = sqlite.connect(source_dbs_path + db_file_name)
        self.c = self.conn.cursor()

    def _init_session(self):
        try:
            self.c.execute('''PRAGMA key = "x'{}'";'''.format(self.key))
            self.c.execute("PRAGMA cipher_compatibility = 3;")
            self.c.execute("PRAGMA kdf_iter = 64000;")
            self.c.execute("PRAGMA cipher_page_size = {};".format(self.cipher_page_size))
            self.c.execute("SELECT name FROM sqlite_master WHERE type='table'")
            # print(self.db_file_name, self.c.fetchall())
        except sqlite.DatabaseError:
            print("db文件：{}，无法解密".format(self.db_file_name))
            return False
        return True

    def exec_sql(self, sql):
        if not self._init_session():
            return

        self.c.execute(sql)
        return self.c.fetchall()

    def exec_update_sql(self, sql):
        if not self._init_session():
            return

        self.c.execute(sql)
        self.conn.commit()

    def decrypt_db(self):
        if not self._init_session():
            return

        decrypt_db_path = "{}{}".format(self.decrypt_db_path, self.db_file_name)
        # if os.path.exists(decrypt_db_path):
        #     os.remove(decrypt_db_path)
        if os.path.exists(decrypt_db_path):
            return

        self.c.execute("PRAGMA cipher_kdf_algorithm = PBKDF2_HMAC_SHA1;")
        self.c.execute("PRAGMA cipher_hmac_algorithm = HMAC_SHA1;")
        self.c.execute("PRAGMA wal_checkpoint;")
        self.c.execute("ATTACH DATABASE '{}' AS plaintext KEY '';".format(decrypt_db_path))
        self.c.execute("SELECT sqlcipher_export('plaintext');")
        self.c.execute("DETACH DATABASE plaintext;")


if __name__ == "__main__":
    db = EncryptDb("../../static/wxdata/", "Backup.db", "../../static/wxdata/decrypt_",
                   "1dda6941f4b2ba3ef923943d1b18789f1584166af5bc6ad99a12b25974bc04cf22a9328944ee8f5d089b42505cca7231",
                   cipher_page_size=4096)
    db.decrypt_db()
